#! /usr/bin/env python
# coding=utf-8

"""
This script processes transcriptome data files (expected in '.genes.results' format) within a given directory, 
performs quantile normalization, and computes z-scores for each gene across samples. The resulting z-scores are 
then saved to a specified output file. This is typically used in RNA-Seq data analysis to standardize gene 
expression levels across different samples or experiments.

Dependencies:
    pandas: for data manipulation and analysis.
    numpy: for numerical operations.
    scipy: specifically, scipy.stats for statistical functions.

Usage:
    Run the script by providing a directory containing '.genes.results' files and an output file path for the z-scores.
    Example command:
    python transcriptome_zscore.py <path_to_directory> <output_file_path>
"""

import os
import pandas as pd
import numpy as np
from scipy.stats import zscore
import shutil
import argparse



def extract_tpm(directory):
    # 获取所有以 .genes.results 结尾的文件名
    files = [file for file in os.listdir(directory) if file.endswith('.genes.results')]
    
    # 初始化一个空的 dataframe 用于存储 TPM 值
    master_df = pd.DataFrame()
    
    for file in files:
        # 读取每一个文件
        df = pd.read_csv(os.path.join(directory, file), sep='\t')
        
        # 从 dataframe 中提取 'gene_id' 和 'TPM' 列
        gene_tpm = df[['gene_id', 'TPM']].copy()
        
        # 设置 gene_id 为索引，并重新命名 'TPM' 列为文件名（但去掉 .genes.results 后缀）
        gene_tpm.set_index('gene_id', inplace=True)
        gene_tpm.rename(columns={'TPM': file.replace('.genes.results', '')}, inplace=True)
        
        # 将这个 dataframe 合并到 master_df
        if master_df.empty:
            master_df = gene_tpm
        else:
            master_df = master_df.join(gene_tpm, how='outer')
            
    return master_df

def quantile_normalize(df):
    """
    This function takes a dataframe and performs quantile normalization on it.
    """
    #print(df)
    # Compute the rank
    df_rank = df.stack().groupby(df.rank(method='first').stack().astype(int)).mean()
    #print(df_rank)
    # Map ranks to mean values
    df_qn = df.rank(method='min').stack().astype(int).map(df_rank).unstack()
    #print(df_qn)
    return df_qn

def quantile_normalize_by_row(df):
    """
    This function takes a dataframe and performs quantile normalization on its rows.
    """
    # Transpose the dataframe
    df_T = df.transpose()
    
    # Compute the rank
    df_rank = df_T.stack().groupby(df_T.rank(method='first').stack().astype(int)).mean()
    
    # Map ranks to mean values
    df_qn_T = df_T.rank(method='min').stack().astype(int).map(df_rank).unstack()
    
    # Transpose the dataframe back
    df_qn = df_qn_T.transpose()
    
    return df_qn
    
def compute_species_zscore(df):
    species_mean = df.mean(axis=1)  # 这里我们计算行的平均值，也就是每个基因across所有samples的平均值
    zscore = (df.subtract(species_mean, axis=0)).div(species_mean, axis=0)
    return zscore

def main(directory, output_prefix):
    tpm_master_df = extract_tpm(directory)
    tpm_master_df = tpm_master_df[tpm_master_df.sum(axis=1) > 1 * len(tpm_master_df.columns)]
    tpm_master_df = tpm_master_df[tpm_master_df.mean(axis=1) > 10]
    log_tpm_master_df = np.log2(tpm_master_df + 1)

    # 保存对数转换后的TPM值到文件

    log_tpm_master_df.to_csv(output_prefix+'_Log2TPM.tsv', sep='\t', index=True, header=True)
    
    log_tpm_master_df_qn = quantile_normalize(log_tpm_master_df)
    log_tpm_master_df_zscore = compute_species_zscore(log_tpm_master_df_qn)
    log_tpm_master_df_zscore.to_csv(output_prefix+'_zscore.tsv', sep='\t', index=True, header=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process transcriptome data to compute zscore.")
    parser.add_argument('directory', type=str, help='Path to the directory containing genes.results files')
    parser.add_argument("output_prefix", help="Prefix for the output file")
    args = parser.parse_args()

    main(args.directory, args.output_prefix)